import logging

from peft import get_peft_model, LoraConfig, PromptEncoderConfig, PromptTuningConfig, PrefixTuningConfig

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class XPeftTuning():
    def __init__(self, config):
        self.config = config

    def getModel(self, model):
        # 微调
        config = self.config
        peft_config = config["peft"]

        if peft_config["user_peft"]:
            tuning_tactics = peft_config["peft_type"]
            if tuning_tactics == 'lora':
                peft_config_param = LoraConfig(r=peft_config["r"], lora_alpha=peft_config["lora_alpha"]
                                               , lora_dropout=peft_config["lora_dropout"],
                                               target_modules=peft_config["target_modules"])

                peft_model = get_peft_model(model, peft_config_param)

                # lora配置会冻结原始模型中的所有层的权重，不允许其反传梯度
                # 但是事实上我们希望最后一个线性层照常训练，只是bert部分被冻结，所以需要手动设置
                for param in peft_model.get_submodule("model").get_submodule("classify").parameters():
                    param.requires_grad = True
                return peft_model

            elif tuning_tactics == "p_tuning":
                peft_config_param = PromptEncoderConfig(task_type="SEQ_CLS",
                                                        num_virtual_tokens=peft_config["num_virtual_tokens"])
                peft_model = get_peft_model(model, peft_config_param)
                return peft_model
            elif tuning_tactics == "prompt_tuning":
                peft_config_param = PromptTuningConfig(task_type="SEQ_CLS",
                                                       num_virtual_tokens=peft_config["num_virtual_tokens"])
                peft_model = get_peft_model(model, peft_config_param)
                return peft_model
            elif tuning_tactics == "prefix_tuning":
                peft_config_param = PrefixTuningConfig(task_type="SEQ_CLS",
                                                       num_virtual_tokens=peft_config["num_virtual_tokens"])
                peft_model = get_peft_model(model, peft_config_param)
                return peft_model
        else:
            return model
